from torch.optim import Adam
import torch
from torch import nn
from tqdm import tqdm
from data import FakeDataset
from models import BertClassifier
from transformers import BertTokenizer
from ignite.metrics import Accuracy,Precision,Recall

def train(model, train_data, val_data, learning_rate, epochs):
  # 通过Dataset类获取训练和验证集
    # DataLoader根据batch_size获取数据，训练时选择打乱样本
    train_dataloader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True, )
    val_dataloader = torch.utils.data.DataLoader(val_data, batch_size=64)
  # 判断是否使用GPU
    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = Adam(model.parameters(), lr=learning_rate)
    m_accuracy, m_precision, m_recall = Accuracy(),Precision(),Recall()

    if use_cuda:
            model = model.cuda()
            criterion = criterion.cuda()
    # 开始进入训练循环
    for epoch_num in range(epochs):
      # 定义两个变量，用于存储训练集的准确率和损失
            total_acc_train = 0
            total_loss_train = 0
      # 进度条函数tqdm
            for train_input, train_label in tqdm(train_dataloader):

                train_label = train_label.to(device)
                mask = train_input['attention_mask'].to(device)
                input_id = train_input['input_ids'].squeeze(1).to(device)

                output = model(input_id, mask)
                # 计算损失
                batch_loss = criterion(output, train_label)
                total_loss_train += batch_loss.item()
                # 计算精度
                acc = (output.argmax(dim=1) == train_label).sum().item()
                total_acc_train += acc
        # 模型更新
                model.zero_grad()
                batch_loss.backward()
                optimizer.step()
            # ------ 验证模型 -----------
            # 定义两个变量，用于存储验证集的准确率和损失
            total_acc_val = 0
            total_loss_val = 0
      # 不需要计算梯度
            with torch.no_grad():
                # 循环获取数据集，并用训练好的模型进行验证
                for val_input, val_label in tqdm(val_dataloader):
          # 如果有GPU，则使用GPU，接下来的操作同训练
                    val_label = val_label.to(device)
                    mask = val_input['attention_mask'].to(device)
                    input_id = val_input['input_ids'].squeeze(1).to(device)
  
                    output = model(input_id, mask)

                    batch_loss = criterion(output, val_label)
                    total_loss_val += batch_loss.item()
                    
                    m_accuracy.update((output.argmax(dim=1) , val_label))
                    m_precision.update((output.argmax(dim=1) , val_label))
                    m_recall.update((output.argmax(dim=1) , val_label))

                    acc = (output.argmax(dim=1) == val_label).sum().item()
                    total_acc_val += acc

            m_accuracy.reset()
            m_precision.reset()
            m_recall.reset()

            print(
                f'''Epochs: {epoch_num + 1} 
              | Train Loss: {total_loss_train / len(train_data): .3f} 
              | Train Accuracy: {total_acc_train / len(train_data): .3f} 
              | Val Loss: {total_loss_val / len(val_data): .3f} 
              | Val Accuracy: {total_acc_val / len(val_data): .3f}''')     

def evaluate(model, test_data):

    test = Dataset(test_data)
    test_dataloader = torch.utils.data.DataLoader(test, batch_size=2)
    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    if use_cuda:
        model = model.cuda()

    total_acc_test = 0
    with torch.no_grad():
        for test_input, test_label in test_dataloader:
              test_label = test_label.to(device)
              mask = test_input['attention_mask'].to(device)
              input_id = test_input['input_ids'].squeeze(1).to(device)
              output = model(input_id, mask)
              acc = (output.argmax(dim=1) == test_label).sum().item()
              total_acc_test += acc   
    print(f'Test Accuracy: {total_acc_test / len(test_data): .3f}')
    
import pandas as pd

if __name__=="__main__":
    MAX_LEN = 500
    EPOCHS = 1
    model = BertClassifier()
    tokenizer = BertTokenizer.from_pretrained('/root/autodl-tmp/bert-base-uncased', do_lower_case=True)
    LR = 1e-6

    df_train = pd.read_csv("/root/autodl-tmp/newsclip/train.tsv", sep='\t', header=0, index_col='index')
    df_test = pd.read_csv("/root/autodl-tmp/newsclip/test.tsv", sep='\t', header=0, index_col='index')
    root_dir = "/root/autodl-tmp/newsclip/origin"

    # 读取数据
    transformed_dataset_train = FakeDataset(df_train, root_dir, tokenizer, MAX_LEN)

    transformed_dataset_val = FakeDataset(df_test, root_dir, tokenizer, MAX_LEN)

    train(model, transformed_dataset_train, transformed_dataset_val, LR, EPOCHS)